#pragma once 
#include <unordered_set>
#include "../common/net.hpp"
#include "../common/message.hpp"

namespace yjz_rpc
{
    namespace server
    {
        class TopicManager
        {
        public:
            using ptr = std::shared_ptr<TopicManager>;
            TopicManager() {}

            void onTopicRequest(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                // 主题创建
                // 主题删除
                // 主题订阅
                // 主题消息发布
                // 主题取消订阅

                TopicOpType topic_optype = msg->opType();
                bool res = true;
                switch (topic_optype)
                {
                    case TopicOpType::TOPIC_CREATE : topicCreate(conn, msg); break;
                    case TopicOpType::TOPIC_REMOVE : topicRemove(conn, msg); break;
                    case TopicOpType::TOPIC_SUBSCRIBE : res = topicSubscribe(conn, msg); break;
                    case TopicOpType::TOPIC_PUBLISH : res = topicPublish(conn, msg); break;
                    case TopicOpType::TOPIC_CANCEL : topicCancel(conn, msg); break;
                    default : return errorResponse(conn, msg, RCode::RCODE_INVALID_OPTYPE);
                }
                if (res == false) 
                {
                    return errorResponse(conn, msg, RCode::RCODE_NOT_FOUND_TOPIC);
                }
                else
                {
                    return topicResponse(conn, msg);
                }
            }

            // 订阅者在连接断开时回调，删除相关的信息
            void onShutdown(const BaseConnection::ptr& conn)
            {
                // 如果断开连接的是发布者，则不需要任何操作；如果是订阅者，则要删除相关信息
                std::vector<Topic::ptr> topics;
                Subscriber::ptr subscriber;
                {
                    // 找到所有受影响的主题对象
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it_sub = _subscribers.find(conn);
                    if (it_sub == _subscribers.end())
                    {
                        return;
                    }
                    subscriber = it_sub->second;
                    for (auto& topic_name : subscriber->topics)
                    {
                        auto it_topic = _topics.find(topic_name);
                        if (it_topic == _topics.end())
                        {
                            continue;
                        }
                        topics.push_back(it_topic->second);
                    }
                    // 从订阅者信息中删除订阅者
                    _subscribers.erase(it_sub);
                }
                // 从受影响的主题对象中，删除订阅者
                for (auto& topic : topics)
                {
                    topic->removeSubscriber(subscriber);
                }
            }
        private:
            void errorResponse(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg, RCode rcode)
            {
                auto rsp_msg = MessageFactory::create<TopicResponse>();
                rsp_msg->setId(msg->getId());
                rsp_msg->setMType(MType::RSP_TOPIC);
                rsp_msg->setRcode(rcode);
                conn->send(rsp_msg);
            }

            void topicResponse(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                auto rsp_msg = MessageFactory::create<TopicResponse>();
                rsp_msg->setId(msg->getId());
                rsp_msg->setMType(MType::RSP_TOPIC);
                rsp_msg->setRcode(RCode::RCODE_OK);
                conn->send(rsp_msg);
            }
        private:
            void topicCreate(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                std::unique_lock<std::mutex> lock(_mutex);
                std::string topic_name = msg->topicKey();
                auto topic = std::make_shared<Topic>(topic_name);
                _topics.insert(std::make_pair(topic_name, topic));
            }

            void topicRemove(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                std::unordered_set<Subscriber::ptr> subs;
                std::string topic_name = msg->topicKey();
                {
                    // 1.首先需要知道当前主题都有哪些订阅者订阅，然后从订阅者中删除该主题信息
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it = _topics.find(topic_name);  // 判断是否有该主题的信息
                    if (it == _topics.end()) 
                    {
                        return;
                    }
                    subs = it->second->subscribers;
                    _topics.erase(it);
                }
                // 2.删除主题名称与主题对象的对应关系
                for (auto& subscriber : subs)
                {
                    subscriber->removeTopic(topic_name);
                }
            }

            bool topicSubscribe(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                // 1.先找到主题和订阅者
                // 如果没找到主题，则报错；如果没找到订阅者，则新建
                Topic::ptr topic;
                Subscriber::ptr subscriber;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it_topic = _topics.find(msg->topicKey());
                    if (it_topic == _topics.end())
                    {
                        return false;
                    }
                    topic = it_topic->second;
                    auto it_sub = _subscribers.find(conn);
                    if (it_sub != _subscribers.end())
                    {
                        subscriber = it_sub->second;
                    }
                    else
                    {
                        subscriber = std::make_shared<Subscriber>(conn);
                        _subscribers.insert(std::make_pair(conn, subscriber));
                    }
                }
                // 2.在主题对象中新增一个订阅者连接，在订阅者对象中新增一个订阅的主题
                topic->appendSubscriber(subscriber);
                subscriber->appendTopic(msg->topicKey());
                return true;
            }

            void topicCancel(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                Topic::ptr topic;
                Subscriber::ptr subscriber;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it_topic = _topics.find(msg->topicKey());
                    if (it_topic != _topics.end())
                    {
                        topic = it_topic->second;
                    }
                    auto it_sub = _subscribers.find(conn);
                    if (it_sub != _subscribers.end())
                    {
                        subscriber = it_sub->second;
                    }
                }
                if (subscriber) subscriber->removeTopic(msg->topicKey());
                if (topic && subscriber) topic->removeSubscriber(subscriber);
            }

            bool topicPublish(const BaseConnection::ptr& conn, const TopicRequest::ptr& msg)
            {
                Topic::ptr topic;
                {
                    std::unique_lock<std::mutex> lock(_mutex);
                    auto it_topic = _topics.find(msg->topicKey());
                    if (it_topic == _topics.end())
                    {
                        return false;
                    }
                    topic = it_topic->second;
                }
                topic->pushMessage(msg);
                return true;
            }
        private:
            struct Subscriber
            {
                using ptr = std::shared_ptr<Subscriber>;
                std::mutex mutex;
                BaseConnection::ptr conn;
                std::unordered_set<std::string> topics; // 当前订阅者都订阅了哪些主题

                Subscriber(const BaseConnection::ptr& c) : conn(c) {}

                void appendTopic(const std::string& topic_name)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    topics.insert(topic_name);
                }

                void removeTopic(const std::string& topic_name)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    topics.erase(topic_name);
                }
            };

            struct Topic
            {
                using ptr = std::shared_ptr<Topic>;
                std::mutex mutex;
                std::string topic_name;
                std::unordered_set<Subscriber::ptr> subscribers; // 当前主题都有哪些连接订阅

                Topic(const std::string& name) : topic_name(name) {}

                void appendSubscriber(const Subscriber::ptr& subscriber) 
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    subscribers.insert(subscriber);
                }

                void removeSubscriber(const Subscriber::ptr& subscriber)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    subscribers.erase(subscriber);
                }

                void pushMessage(const BaseMessage::ptr& msg)
                {
                    std::unique_lock<std::mutex> lock(mutex);
                    for (auto& subscriber : subscribers)
                    {
                        subscriber->conn->send(msg);
                    }
                }
            };
        private:
            std::mutex _mutex;
            std::unordered_map<std::string, Topic::ptr> _topics; // 主题名称和主题的对应关系
            std::unordered_map<BaseConnection::ptr, Subscriber::ptr> _subscribers;  // 连接和订阅者的对应关系
        };
    }
}